from datetime import datetime
from multiprocessing import cpu_count, Pool
import akshare as ak
import pandas as pd

from program.Config import *


def index_zh_a_hist(
        symbol: str = "000001",
        period: str = "daily",
        start_date: str = "19700101",
        end_date: str = "22220101",
) -> pd.DataFrame:
    symbol = symbol.replace("sz", "").replace("sh", "")
    return ak.index_zh_a_hist(symbol, period, start_date, end_date)


def update_single_index(index_code: str):
    to_csv_path = index_data_path + '%s.csv' % index_code
    df_new = pd.DataFrame()
    if os.path.exists(to_csv_path):  # 如果文件存在就接着更新
        df_old = pd.read_csv(to_csv_path, encoding='gbk')
        df_old['candle_end_time'] = pd.to_datetime(df_old['candle_end_time'])
        # 虽然原定的是一周数据, 但是还是用条件筛选一下
        # 获取旧数据最后一天的日期
        last_date = df_old.iloc[-1]['candle_end_time']
        df_new = index_zh_a_hist(symbol=index_code, start_date=last_date.strftime('%Y%m%d'))
        # 转换字段名字
        df_new.rename(columns={'日期': 'candle_end_time',
                               '开盘': 'open',
                               '最高': 'high',
                               '最低': 'low',
                               '收盘': 'close',
                               '成交额': 'amount'}, inplace=True)
        # 转换日期格式
        df_new['candle_end_time'] = pd.to_datetime(df_new['candle_end_time'])
        df_new = df_new[['candle_end_time', 'open', 'high', 'low', 'close', 'amount']]
        df_new['info'] = '{}'
        # 获取新数据中大于旧数据最后一天的数据
        df_new = df_new[df_new['candle_end_time'] > last_date]
        if df_new.shape[0] == 0:
            print(f"{index_code} 指数没有新数据需要更新!")
            return
        # 合并新旧数据
        df_new = pd.concat([df_old, df_new], ignore_index=True)
    else:
        df_new = index_zh_a_hist(symbol=index_code)
        # 转换字段名字
        df_new.rename(columns={'日期': 'candle_end_time',
                               '开盘': 'open',
                               '最高': 'high',
                               '最低': 'low',
                               '收盘': 'close',
                               '成交额': 'amount'}, inplace=True)
        # 转换日期格式
        df_new['candle_end_time'] = pd.to_datetime(df_new['candle_end_time'])
        df_new = df_new[['candle_end_time', 'open', 'high', 'low', 'close', 'amount']]
    # 保存数据
    df_new.to_csv(to_csv_path, encoding='gbk', index=False)


def update_index():
    # 标记开始时间
    start_time = datetime.now()
    # 获取指数列表
    index_list = os.listdir(index_data_path)
    # 遍历index_list，去掉.csv后缀
    index_list = [index.replace('.csv', '') for index in index_list]
    # 遍历指数列表，更新数据
    # for index_code in index_list:
    #     update_single_index(index_code)
    # 并行提速的办法
    with Pool(max(cpu_count() - 1, 1)) as pool:
        pool.map(update_single_index, sorted(index_list))
    print('更新指数数据完毕，消耗时间', datetime.now() - start_time)

# 下载数据并更新：https://www.quantclass.cn/data/stock/stock-trading-data-pro@recent7d
def update_stock_daily():
    # 标记开始时间
    start_time = datetime.now()
    stock_file_list = os.listdir(stock_update_data_path)
    # 并行提速的办法
    with Pool(max(cpu_count() - 1, 1)) as pool:
        pool.map(update_single_stock_daily, sorted(stock_file_list))
    print('更新日线数据完毕，消耗时间', datetime.now() - start_time)


def update_single_stock_daily(file_name: str):
    new_csv_path = stock_update_data_path + file_name
    to_csv_path = stock_data_path + '/' + file_name
    df_new = pd.read_csv(new_csv_path, encoding='gbk', skiprows=1)
    df_new['交易日期'] = pd.to_datetime(df_new['交易日期'])

    # 兼容之前的字段
    if '申万一级行业名称' not in df_new.columns:
        df_new['申万一级行业名称'] = df_new['新版申万一级行业名称']
    if '申万二级行业名称' not in df_new.columns:
        df_new['申万二级行业名称'] = df_new['新版申万二级行业名称']
    if '申万三级行业名称' not in df_new.columns:
        df_new['申万三级行业名称'] = df_new['新版申万三级行业名称']
    # 调整字段顺序
    df_new = df_new[
        ['股票代码', '股票名称', '交易日期', '开盘价', '最高价', '最低价', '收盘价', '前收盘价', '成交量',
         '成交额', '流通市值', '总市值', '净利润TTM', '现金流TTM', '净资产', '总资产', '总负债', '净利润(当季)',
         '中户资金买入额', '中户资金卖出额', '大户资金买入额', '大户资金卖出额', '散户资金买入额',
         '散户资金卖出额', '机构资金买入额', '机构资金卖出额', '沪深300成分股', '上证50成分股', '中证500成分股',
         '中证1000成分股', '创业板指成分股', '申万一级行业名称', '申万二级行业名称', '申万三级行业名称',
         '新版申万一级行业名称', '新版申万二级行业名称', '新版申万三级行业名称', '09:35收盘价', '09:45收盘价',
         '09:55收盘价']]

    try:
        if file_name in os.listdir(stock_data_path):
            df_old = pd.read_csv(to_csv_path, encoding='gbk', skiprows=1)
            df_old['交易日期'] = pd.to_datetime(df_old['交易日期'])
            # 虽然原定的是一周数据, 但是还是用条件筛选一下
            # 获取旧数据最后一天的日期
            last_day = df_old.iloc[-1]['交易日期']
            # 获取大于这个日期的新数据
            df_new = df_new[df_new['交易日期'] > pd.to_datetime(last_day)]
            # df_new为空则跳过
            if df_new.shape[0] == 0:
                print(f"{file_name} 股票没有新数据需要更新!")
                return
            # df_new 追加到csv
            df_new.to_csv(to_csv_path, encoding='gbk', index=False, mode='a', header=False)
            print(f"更新 {file_name} 股票数据完成!")
        else:
            # 如果不在之前的列表说明是新股，直接保存
            # 先将头文件输出
            pd.DataFrame(columns=['本数据策略分享会专用，由邢不行整理，微信：xbx6660']).to_csv(to_csv_path, index=False,
                                                                                            encoding='gbk')
            df_new.to_csv(to_csv_path, encoding='gbk', mode='a', index=False)
            print(f"添加 {file_name} 股票数据完成!")
    except Exception as e:
        print("\033[31m更新 %s 股票数据失败，原因：%s\030m" % (file_name, e))
        return


def delete_stock_record():
    # 标记开始时间
    start_time = datetime.now()
    stock_file_list = os.listdir(stock_update_data_path)
    # 并行提速的办法
    with Pool(max(cpu_count() - 1, 1)) as pool:
        pool.map(delete_single_stock_record, sorted(stock_file_list))
    print('删除日线数据完毕，消耗时间', datetime.now() - start_time)

def delete_single_stock_record(file_name: str):
    from_date = '2023-11-27'
    to_csv_path = stock_data_path + '/' + file_name
    if os.path.exists(to_csv_path):
        df_old = pd.read_csv(to_csv_path, encoding='gbk', skiprows=1)
        df_old['交易日期'] = pd.to_datetime(df_old['交易日期'])
        df_old = df_old[df_old['交易日期'] < pd.to_datetime(from_date)]
        # 先将头文件输出
        pd.DataFrame(columns=['本数据策略分享会专用，由邢不行整理，微信：xbx6660']).to_csv(to_csv_path, index=False,
                                                                                        encoding='gbk')
        df_old.to_csv(to_csv_path, encoding='gbk', mode='a', index=False)
        print(f"删除 {file_name} 股票{from_date}之后的数据完成!")

def update_stock_financial():
    start_time = datetime.now()
    stock_file_list = os.listdir(finance_update_data_path)
    # 并行提速的办法
    with Pool(max(cpu_count() - 1, 1)) as pool:
        pool.map(update_single_stock_financial, sorted(stock_file_list))
    print('更新财务数据完毕，消耗时间', datetime.now() - start_time)


def update_single_stock_financial(file_name: str):
    # 之前的财务数据路径
    to_csv_path = finance_data_path + '/' + file_name
    # 更新的财务数据路径
    new_csv_path = finance_update_data_path + file_name
    # 获取更新的文件
    new_file = os.listdir(new_csv_path)[0]
    df_new = pd.read_csv(f'{new_csv_path}/{new_file}', skiprows=1, encoding='gbk', parse_dates=['抓取时间'])
    # 判断该股票是否存在财务数据
    if os.path.exists(to_csv_path):
        # 获取该文件夹下的文件
        file = os.listdir(to_csv_path)[0]
        # 读取财务数据
        finance_df = pd.read_csv(f'{to_csv_path}/{file}', skiprows=1, encoding='gbk', parse_dates=['抓取时间'])
        # 获取旧数据最后抓取时间
        last_day = finance_df.iloc[-1]['抓取时间']
        # 获取大于这个日期的新数据
        df_new = df_new[df_new['抓取时间'] > pd.to_datetime(last_day)]
        # df_new为空则跳过
        if df_new.shape[0] == 0:
            print(f"{file_name} 股票没有新财务数据需要更新!")
            return
        # 抓取时间排序
        df_new.sort_values('抓取时间', inplace=True)
        # df_new 追加到csv
        df_new.to_csv(f'{to_csv_path}/{file}', encoding='gbk', index=False, mode='a', header=False)
        print(f"更新 {file_name} 财务数据完成!")
    else:
        # 创建文件夹
        os.mkdir(to_csv_path)
        # 输出文件头
        pd.DataFrame(columns=['数据由邢不行整理，对数据字段有疑问的，可以直接微信私信邢不行，微信号：xbx6660']).to_csv(f'{to_csv_path}/{new_file}', index=False, encoding='gbk')
        # 抓取时间排序
        df_new.sort_values('抓取时间', inplace=True)
        # 输出数据
        df_new.to_csv(f'{to_csv_path}/{new_file}', encoding='gbk', mode='a', index=False)
        print(f"添加 {file_name} 财务数据完成!")



def check_stock_daily_data():
    stock_file_list = os.listdir(stock_data_path)
    for file_name in stock_file_list:
        try:
            df = pd.read_csv(stock_data_path + '/' + file_name, encoding='gbk', skiprows=1)
            df['交易日期'] = pd.to_datetime(df['交易日期'])
        except Exception as e:
            print(f"文件 {file_name} 不能转化为dataframe")
            print(e)


if __name__ == '__main__':
    # delete_stock_record()
    # update_single_stock_daily('bj430300.csv')
    # delete_single_stock_record('bj430300.csv')
    update_stock_daily()
    update_index()
    # update_stock_financial() # 财务数据一个月更新一次
    # update_single_index('sh000300')